#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 使用前需运行：sudo chmod 777 /dev/ttyACM0 以确保串口权限
# Z轴做一次启动去皮

import time
from collections import deque

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import WrenchStamped
from std_msgs.msg import Header
import serial

# ==== Modbus 指令（逐轴读取）====
hex_data1 = bytes([0xFE, 0x04, 0x00, 0x00, 0x00, 0x01, 0x25, 0xC5])  # X轴
hex_data2 = bytes([0xFE, 0x04, 0x00, 0x01, 0x00, 0x01, 0x74, 0x05])  # Y轴
hex_data3 = bytes([0xFE, 0x04, 0x00, 0x02, 0x00, 0x01, 0x84, 0x05])  # Z轴
hex_datas = (hex_data1, hex_data2, hex_data3)


class ForceSensorReader(Node):
    def __init__(self):
        super().__init__('force_sensor_reader')

        # ===== 参数 =====
        self.serial_port            = self.declare_parameter('serial_port', '/dev/ttyACM0').value
        self.baudrate               = int(self.declare_parameter('baudrate', 38400).value)
        self.timeout_s              = float(self.declare_parameter('timeout_s', 0.2).value)

        self.ema_alpha              = float(self.declare_parameter('ema_alpha', 0.3).value)
        self.jump_threshold         = float(self.declare_parameter('jump_threshold', 5.0).value)
        self.median_window          = int(self.declare_parameter('median_window', 5).value)
        self.deadband               = float(self.declare_parameter('deadband', 0.01).value)
        self.sensor_frame           = self.declare_parameter('sensor_frame', 'force_sensor').value
        self.log_every_n            = int(self.declare_parameter('log_every_n', 10).value)

        self._log_count = 0

        # ===== 串口 =====
        self.ser = serial.Serial(port=self.serial_port, baudrate=self.baudrate, timeout=self.timeout_s)
        self.ser.reset_input_buffer()
        self.ser.reset_output_buffer()

        # ===== 发布者 =====
        self.publisher_ = self.create_publisher(WrenchStamped, '/force_data', 10)

        # 状态缓存
        self._last_good = {'x': 0.0, 'y': 0.0, 'z': 0.0}
        self._med_buf = {
            'x': deque(maxlen=self.median_window),
            'y': deque(maxlen=self.median_window),
            'z': deque(maxlen=self.median_window),
        }

        # Z轴去皮偏置
        self._z_offset = None
        self._tare_attempts = 10  # 取多次平均值更稳

        # 100Hz
        self.timer = self.create_timer(0.01, self.read_force_data)

    # ============ Modbus 工具 ============
    def _modbus_crc16(self, data: bytes) -> int:
        crc = 0xFFFF
        for b in data:
            crc ^= b
            for _ in range(8):
                if crc & 0x0001:
                    crc = (crc >> 1) ^ 0xA001
                else:
                    crc >>= 1
        return crc

    def _write_and_read_frame(self, req: bytes, resp_len: int = 7, retries: int = 3) -> bytes:
        for _ in range(retries):
            try:
                self.ser.reset_input_buffer()
                self.ser.write(req)
                buf = b""
                t0 = time.time()
                while len(buf) < resp_len and (time.time() - t0) < self.ser.timeout:
                    need = resp_len - len(buf)
                    chunk = self.ser.read(need)
                    if chunk:
                        buf += chunk
                    else:
                        time.sleep(0.001)

                if len(buf) != resp_len:
                    continue

                frame_wo_crc = buf[:-2]
                crc_calc = self._modbus_crc16(frame_wo_crc)
                crc_lo = crc_calc & 0xFF
                crc_hi = (crc_calc >> 8) & 0xFF
                if buf[-2] != crc_lo or buf[-1] != crc_hi:
                    self.get_logger().warn("CRC mismatch, retrying...")
                    continue

                if buf[2] != 0x02:
                    self.get_logger().warn(f"Unexpected byte count {buf[2]}, retrying...")
                    continue

                return buf
            except Exception as e:
                self.get_logger().warn(f"Frame read error: {e}; retrying...")

        raise TimeoutError("Modbus frame read failed after retries.")

    # ============ 主回调 ============
    def read_force_data(self):
        header = Header()
        header.stamp = self.get_clock().now().to_msg()
        header.frame_id = self.sensor_frame

        try:
            axis_map = {0: 'x', 1: 'y', 2: 'z'}
            vals = [0.0, 0.0, 0.0]

            for i in range(3):
                key = axis_map[i]
                try:
                    frame = self._write_and_read_frame(hex_datas[i], resp_len=7, retries=3)
                    val = (frame[3] << 8) | frame[4]
                    if val > 32767:
                        val -= 65536
                    raw = val / 100.0
                except Exception as e:
                    self.get_logger().warn(f"Axis {i} read failed: {e}")
                    raw = self._last_good[key]

                # Z轴去皮逻辑
                if key == 'z':
                    if self._z_offset is None:
                        # 第一次启动时取多次平均作为偏置
                        if not hasattr(self, "_z_samples"):
                            self._z_samples = []
                        self._z_samples.append(raw)
                        if len(self._z_samples) >= self._tare_attempts:
                            self._z_offset = sum(self._z_samples) / len(self._z_samples)
                            self.get_logger().info(f"Z-axis tare complete, offset={self._z_offset:.3f}")
                        raw = 0.0  # 去皮未完成时发布0
                    else:
                        raw = raw - self._z_offset

                # —— 抗跳变滤波 —— 
                mbuf = self._med_buf[key]
                mbuf.append(raw)
                med = sorted(mbuf)[len(mbuf)//2]

                prev = self._last_good[key]
                if abs(med - prev) > self.jump_threshold:
                    med = prev + 0.5 * (med - prev)

                smoothed = (1.0 - self.ema_alpha) * prev + self.ema_alpha * med

                if abs(smoothed) < self.deadband:
                    smoothed = 0.0

                self._last_good[key] = smoothed
                vals[i] = smoothed

            fx, fy, fz = vals

            msg = WrenchStamped()
            msg.header = header
            msg.wrench.force.x = fx
            msg.wrench.force.y = fy
            msg.wrench.force.z = fz
            self.publisher_.publish(msg)

            self._log_count += 1
            if self._log_count % self.log_every_n == 0:
                self.get_logger().info(f"RAW(filtered) X={fx:.3f}, Y={fy:.3f}, Z={fz:.3f}")

        except Exception as e:
            self.get_logger().error(f"Error reading force data: {e}")


def main(args=None):
    rclpy.init(args=args)
    node = ForceSensorReader()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
